import numpy as np
import time
import os
import scipy.io
from scipy.optimize import curve_fit
import matplotlib.pyplot as plt
import json
import h5py

debug_mode = False

global channel_return

class DataDoc():
    def __init__(self, disk_path=r'D:\Data', title="", path=None, debug_mode=False):
        if path==None:
            self.path_prefix = disk_path + time.strftime("\%Y\%Y%m\%Y%m%d\\")
        else:
            self.path_prefix = path
        #if not os.path.exists(self.path_prefix) and not debug_mode:
        path_prefix = self.path_prefix
        if not os.path.exists(self.path_prefix):
            print("Data will be save into ",self.path_prefix)
            os.makedirs(self.path_prefix)
        
        self.title            = title
        self.xlabel           = "Duration"
        self.ylabel           = "Counts"
        self.SystemParameters = {"CONTENT":"EMPTY"}
        self.IonParameters    = {"CONTENT":"EMPTY"}
        self.Sequence         = {"CONTENT":"EMPTY"}
        self.RawData          = {"Duration":np.array([],dtype=np.float32),"Counts":np.array([],dtype=np.float32)}
        self.DataFileName     = ""

    def new_data_file(self):
        file_name = self.title + time.strftime("-%Y%m%d%H%M%S")
        return file_name

    def jsonify(self, data):
        data_return = list(data)
        return data_return

    def reset_raw_data(self, xtype="Duration"):
        self.xlabel = xtype
        self.RawData          = {self.xlabel: np.array([],dtype=np.float32), self.ylabel: []}
        self.DataFileName     = self.new_data_file()
        print(self.path_prefix+self.DataFileName+".hdf5")

    def append_raw_data(self, x, y):
        self.RawData[self.xlabel] = np.append(self.RawData[self.xlabel],x)
        self.RawData[self.ylabel] = self.RawData[self.ylabel] + [y]

    def generate_json(self):
        file_name = self.DataFileName
        json_name = file_name + ".json"
        json_path = self.path_prefix + json_name
        hdf5_name = file_name + ".hdf5"
        hdf5_path = self.path_prefix + hdf5_name

        self.DataFileName = hdf5_path

        data_to_save = self.__dict__.copy()
        del data_to_save["RawData"]
        with open(json_path, 'w', encoding='utf-8') as f:
            json.dump(data_to_save, f, ensure_ascii=False, indent=4)

        with h5py.File(hdf5_path, 'w') as f:
             f[self.xlabel] = self.RawData[self.xlabel]
             f[self.ylabel] = np.array(self.RawData[self.ylabel],dtype=np.float32)
        return file_name


class scanParameterType():
    time = 1
    frequency = 2

def raw_count(path):
    '''
    return two lists: the first is the list of scan parameter(time or frequency), the second is the related list of matrices
    '''

    #  try:
    #      mat = h5py.File(path, 'r')
    # except Exception as e:
    #     if type(e) == FileNotFoundError:
    #        print("Error")
            #path = path_prefix + path
            #mat = h5py.File(path, 'r')
    mat = h5py.File(path, 'r')
    ylabel, xlabel = mat.keys()
    list_time = []
    list_values = []
    for i in range(len(mat[xlabel])):
        timeStr = mat[xlabel][i]
        list_time.append(timeStr)
        list_values.append(mat[ylabel][i])
    return list_time, list_values


def verify(total_channels, show_channels):
    assert type(total_channels) == int
    assert (total_channels > 0)
    if type(show_channels) == int:
        show_channels = [show_channels]
    if show_channels == None:
        show_channels = [i for i in range(total_channels)]
    assert type(show_channels) == list
    return total_channels, show_channels

def average(filename, threshold=None):
    total_channels=4
    if threshold!=None:
        assert type(threshold) == list
        list_time, list_values = raw_count(filename)
        repeat = len(list_values[0]) // total_channels
        for i in range(len(list_values)):
            list_values[i] = np.transpose(list_values[i].reshape(repeat, total_channels))
            for j in range(total_channels):
                list_values[i][j] = np.array([np.where(list_values[i][j] > threshold[j], 1, 0)])
        for i in range(len(list_values)):
            list_values[i] = np.mean(list_values[i], axis=1)
                
        #list_values = np.transpose(list_values)
        #list_values = np.array([np.where(list_values[i] > threshold[i], 1, 0) for i in np.arange(len(list_values))])
        #list_values = np.transpose(list_values)
        #list_values = [np.mean(value,axis=0) for value in list_values]
    else:
        list_time, list_values = raw_count(filename)
        repeat = len(list_values[0]) // total_channels
        for i in range(len(list_values)):
            list_values[i] = np.transpose(list_values[i].reshape(repeat, total_channels))
            for j in range(total_channels):
                list_values[i][j] = np.mean(list_values[i][j])
        for i in range(len(list_values)):
            list_values[i] = np.mean(list_values[i], axis=1)
    return [np.array(list_time), np.array(list_values)]

def average_plot(filename, threshold=None, show_channels=[3]):
    total_channels=4
    total_channels, show_channels = verify(total_channels, show_channels)
    list_time, list_values = average(filename, threshold=threshold)
    list_values = np.transpose(list_values)
    ch_number = len(show_channels)
    plt.figure(figsize=(8,8))
    for i in range(ch_number):
        ch_idx = show_channels[i]
        plt.subplot(ch_number,1,i+1)
        plt.plot(list_time, list_values[ch_idx])
        plt.title('This is channel {}'.format(ch_idx))
        plt.xlabel('time '+r'$(\mu s)$')
        plt.ylabel('average count')
        plt.tight_layout()
    return

def pre_process(list_matrices, convert_matrix, threshold):
    results = []
    total_channels = 4
    repeat = len(list_matrices[0])//total_channels
    for index, raw_matrix in enumerate(list_matrices):
        new_matrix = raw_matrix.reshape(repeat, total_channels)
        new_matrix = np.transpose(new_matrix)
        new_matrix = convert_matrix @ new_matrix
        if threshold != None:
            for j in range(new_matrix.shape[0]):
                new_matrix[j] = np.where(new_matrix[j]>threshold[j], 1, 0)
        avrg = np.sum(new_matrix, axis = 1)/repeat
        results.append(avrg)

    return results

def average_fit(file_name, convert_matrix = None, threshold = None, para_type = scanParameterType.time):
    assert(type(convert_matrix) != type(None))

    list_para, list_matrices = raw_count(file_name)
    avrg_data = pre_process(list_matrices, convert_matrix)
    results = avrg_data
    if threshold != None:
        results = []
        for avrg in avrg_data:
            avrg = (avrg > threshold).astype(int)
            results.append(avrg)
    return results

def histogram_plot(bright_data, dark_data):
    max_index = max(max(bright_data),max(dark_data))
    hist_bright = np.histogram(bright_data, bins=np.arange(max_index), density=True)
    hist_dark = np.histogram(dark_data, bins=np.arange(max_index), density=True)
    plt.figure()
    plt.bar(hist_bright[1][:-1]+0.5,hist_bright[0])
    plt.bar(hist_dark[1][:-1]+0.5,hist_dark[0])
    plt.show() 
    fidelity_list = [(sum(hist_bright[0][i:])+sum(hist_dark[0][:i]))/2 for i in range(max_index)]
    threshold = max(range(len(fidelity_list)), key=fidelity_list.__getitem__)
    return [threshold + 1, fidelity_list[threshold]]

def correlation(file_name, convert_matrix = None, threshold = None, para_type = scanParameterType.time):
    pass


def cosine_func(x,a0,a1,a2,a3):
    return a0 * np.sin(a1*x+a2) + a3

def gaussian_func(x,a,mu,sigma):
    return a*np.exp(-(x-mu)**2/(2*sigma**2))

def gaussian_func2(x, a, mu, sigma_reverse):
    return a*np.exp(-(x-mu)**2 * sigma_reverse**2)

def thermal_single_func(x, p0, gamma, omega):
    return 1/2*p0*(1 - np.exp(-gamma*x)*np.cos(omega*x))

def combinatorial_number(n,m):
    return math.factorial(n) // (math.factorial(m)*math.factorial(n-m))

def Laguerre(n,x):
    sum = 0
    for k in range(n+1):
        sum += (-1)**k*combinatorial_number(n+1,n-k)*(x**k / math.factorial(k))
    return sum

def thermal_func(x, *args):
    '''
    pn, gamma, omega are all lists
    Laguerre function:
        L_n^a(x) = \sum_{k=0}^{n} (-1)^k C_{n+a}^{n-k} x^k/k!
    '''
    eta = 0.098 #eta is a pre-given constant

    n = len(args) // 2

    assert (len(args) == 2*n+1)

    pn = np.array(args[0:n])
    gamma = np.array(args[n:2*n])
    omega = args[-1]


    omega_l = np.array([omega]+[0 for i in range(n-1)])
    for i in range(1,n):
        omega_l[i] = omega * np.exp(-eta*eta/2) * eta * np.sqrt(1/(i+1)) * Laguerre(i,eta*eta)
    sum_p = 0
    for i in range(n):
        sum_p += 1/2*pn[i]*(1 - np.exp(-gamma[i]*x) * np.cos(omega_l[i]*x))
    return sum_p

def automatic_find_initial_omega(xdata, ydata):
    pass

def check_fitting_quality(ion, xdata, ydata, y_fit):
    pass

def gaussian_fit(fileName, convert_matrix = None, threshold = None, para_type = scanParameterType.frequency, plot_figure = False):
    list_frequency, list_matrices = raw_count(fileName)
    avrg_data_all = pre_process(list_matrices, convert_matrix, threshold)

    ion_number = convert_matrix.shape[0]
    fit_paras = []

    for ion_index in range(ion_number):
        
        avrg_single_ion = [avrg[ion_index] for avrg in avrg_data_all]
        xdata = np.array(list_frequency)
        ydata = np.array(avrg_single_ion)

        #mean,std=scipy.stats.norm.fit(ydata)
        a0 = max(ydata)
        
        if a0 == 0:
            fit_paras.append([0,0,1])
            continue
        a1 = xdata[np.argmax(ydata)]
        #a2 = np.std(ydata)
        a2 = np.std(ydata) * (xdata[1] - xdata[0]) / a0
        p0 = [a0, a1, a2]

        #a2 = sum(y * (x - a1)**2)
        #sigma_reverse = 1/(a2 * np.sqrt(2)
        #p0 = [a0, a1, sigma_reverse]
        #p_l = [a0/2, xdata[0], a2/2]
        #p_h = [a0*2, xdata[-1], a2*2]

        #print(p0)
        popt, pcov = curve_fit(gaussian_func, xdata, ydata, p0=p0)
        #popt, pcov = curve_fit(gaussian_func2, xdata, ydata, p0=p0)
        fit_paras.append(popt)
        #print(popt)

        fit_data = gaussian_func(xdata, *popt)
        check_fitting_quality(ion_index, xdata, ydata, fit_data)
        #print('fit_paras', popt)

    if plot_figure:
        plt.figure(figsize=(8,8))
        for ion_index in range(ion_number):
            if ion_index != 2:
                continue ##目前只需要第3个通道
            avrg_single_ion = [avrg[ion_index] for avrg in avrg_data_all]
            x_fit = np.linspace(min(list_frequency),max(list_frequency), 100)
            avrg_fit = [gaussian_func(x, *fit_paras[ion_index]) for x in x_fit]

            plt.subplot(ion_number,1,ion_index+1)
            plt.plot(list_frequency, avrg_single_ion)

            xdata = np.array(list_frequency)
            ydata = np.array(avrg_single_ion)
            #a0 = max(ydata)
            #a1 = xdata[np.argmax(ydata)]
            #a2 = sum(y * (x - a1)**2)
            #ydata2 = gaussian_func(xdata, a0, a1, a2)
            #plt.plot(xdata, ydata2)

            plt.plot(x_fit, avrg_fit)
            #print(fit_paras[ion_index])
            plt.title(('This is channel {}, '+r'$\mu $'+'= {:.4f}, '+r'$\sigma = {:.4f}$').format(ion_index, fit_paras[ion_index][1], fit_paras[ion_index][2]))
            #plt.title(('This is ion {}, '+r'$\mu $'+'= {:.4f}, '+r'$\sigma = {:.4f}$').format(ion_index, fit_paras[ion_index][1], np.sqrt(2)/fit_paras[ion_index][2]))
            plt.xlabel('frequency '+' (MHz)')
            plt.ylabel('average count')

        plt.tight_layout()

    return fit_paras


def rabi_fit(fileName, convert_matrix = None, threshold = None, para_type = scanParameterType.frequency, plot_figure = False):
    list_time, list_matrices = raw_count(fileName)
    avrg_data_all = pre_process(list_matrices, convert_matrix, threshold)
    #print(list_matrices)
    #print(avrg_data_all)

    ion_number = convert_matrix.shape[0]

    fit_paras = []

    for ion_index in range(ion_number):
        avrg_single_ion = [avrg[ion_index] for avrg in avrg_data_all]
        fs = np.fft.fftfreq(len(list_time))
        fs = np.fft.fftfreq(len(list_time), list_time[1]-list_time[0])
        Y = abs(np.fft.fft(avrg_single_ion))
        freq = abs(fs[np.argmax(Y[1:])+1])
        #print(freq)
        xdata = np.array(list_time)
        ydata = np.array(avrg_single_ion)
        a0 = (max(avrg_single_ion) - min(avrg_single_ion)) / 2
        a1 = 2 * np.pi * freq
        a2 = 0
        a3 = np.mean(avrg_single_ion)
        p0 = [a0,a1,a2,a3]
        if ion_index != 2:
            fit_paras.append(np.array(p0))
            continue
        popt, pcov = curve_fit(cosine_func, list_time, avrg_single_ion, p0=p0)
        fit_paras.append(popt)

        #print(popt)
        fit_data = cosine_func(xdata, *popt)
        check_fitting_quality(ion_index, xdata, ydata, fit_data)

    if plot_figure:
        plt.figure(figsize=(8,8))
        #print(fit_paras)
        for ion_index in range(ion_number):
            if ion_index != 2:
                continue
            avrg_single_ion = [avrg[ion_index] for avrg in avrg_data_all]
            x_fit = np.linspace(min(list_time),max(list_time), 100)
            avrg_fit = [cosine_func(x, *fit_paras[ion_index]) for x in x_fit]

            plt.subplot(ion_number,1,ion_index+1)
            plt.plot(list_time, avrg_single_ion)
            plt.plot(x_fit, avrg_fit)
            plt.title('This is channel {}, the pi-pulse period is {:.4f} '.format(ion_index, (np.pi)/popt[1])+r'$\mu s$')
            plt.xlabel('time '+r'$(\mu s)$')
            plt.ylabel('average count')
        plt.tight_layout()

    fit_results = [None, None, (np.pi)/fit_paras[2][1], None]
    #fit_results = [(2*np.pi)*1e6/a[1] for a in fit_paras]
    return fit_results

def thermal_fit(fileName, convert_matrix = None, threshold = None, para_type = scanParameterType.frequency, plot_figure = False):
    list_time, list_matrices = raw_count(fileName)
    avrg_data_all = pre_process(list_matrices, convert_matrix)

    ion_number = convert_matrix.shape[0]

    fit_paras = []

    for ion_index in range(ion_number):
        avrg_single_ion = [avrg[ion_index] for avrg in avrg_data_all]
        fs = np.fft.fftfreq(len(list_time))
        fs = np.fft.fftfreq(len(list_time), list_time[1]-list_time[0])
        Y = abs(np.fft.fft(avrg_single_ion))
        freq = abs(fs[np.argmax(Y[1:])+1])
        print(freq)

